/*********************************************************
* unit:    quantise         release 0.33                 *
* purpose: Fast palette quantiser.                       *
* licency:     GPL or LGPL                               *
* Copyright: (c) 2024-2025 Jaroslav Fojtik               *
**********************************************************/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "csext.h"	// InterlockedIncrement()

#include "raster.h"
#include "ras_prot.h"

#include "quantise.h"
#include "OmpFix.h"


inline int32_t sqr(const int32_t val)
{
  return val*val;
}


//#define _DEBUG

///////////////////////////////////////////////////

typedef struct
{
  uint32_t Hits;  
  uint8_t R,G,B;
} TQuantumRGB;


///////////////////////////////////////////////////


void InsertSort(TQuantumRGB *Arr, unsigned indices, uint32_t hits, unsigned x, unsigned y, unsigned z)
{
  for(unsigned i=0; i<indices; i++)
  {
    if(hits > Arr[i].Hits)
    {
      if(i+1 < indices)
        memmove(&Arr[i+1], &Arr[i], indices-i-1);		// overlapping regions.
      Arr[i].Hits = hits;
      Arr[i].R = x;
      Arr[i].G = y;
      Arr[i].B = z;
      break;
    }
  }
}


///////////////////////////////////////////////////


APalette *FindIndices(Raster2DAbstract *RasterRGB, unsigned indices)
{
  if(RasterRGB==NULL || RasterRGB->Size1D<=0 || RasterRGB->Size2D<=0 || indices<=1 || indices>=65535) return NULL;

  Raster3D_32Bit OctTree;
  OctTree.Allocate3D(16,16,16);
  if(OctTree.Data3D==NULL) return NULL;

  TQuantumRGB *Arr = (TQuantumRGB*)calloc(indices,sizeof(TQuantumRGB));
  if(Arr==NULL) return NULL;

  for(unsigned z=0; z<OctTree.Size3D; z++)
  {
    for(unsigned y=0; y<OctTree.Size2D; y++)
        memset(OctTree.GetRow(y,z), 0, 16*4);
  }

#pragma omp parallel
  {
    Raster1DAbstract * const RasRGBrow 
	  =  (RasterRGB->Channels()==3) ? CreateRaster1DRGB(0,RasterRGB->GetPlanes()/3) : CreateRaster1DRGBA(0,RasterRGB->GetPlanes()/4);
    RasRGBrow->Shadow = true;
    RasRGBrow->Size1D = RasterRGB->Size1D;
#pragma omp for ordered schedule(dynamic)
    for(UNS_OMP_IT y=0; y<RasterRGB->Size2D; y++)
    {
      RasRGBrow->Data1D = RasterRGB->GetRow(y);
      if(RasRGBrow->Data1D==NULL) continue;
      for(unsigned x=0; x<RasterRGB->Size1D; x++)
      {
        RGBQuad RGB;
        RasRGBrow->Get(x,&RGB);
        uint32_t *ptr = RGB.R/16 + (uint32_t*)OctTree.GetRow(RGB.G/16,RGB.B/16);
#if defined _OPENMP
        InterlockedIncrement((long*)ptr);
#else
        *ptr++;
#endif
      }
    }

    delete RasRGBrow;
  }

	// Pick up maximal occurancy chunks for a palette.
 for(unsigned z=0; z<OctTree.Size3D; z++)
  for(unsigned y=0; y<OctTree.Size2D; y++)
  {
    const uint32_t * const RGBrow = (uint32_t*)OctTree.GetRow(y,z);
    for(unsigned x=0; x<OctTree.Size1D; x++)
    {
      InsertSort(Arr, indices, RGBrow[x], x, y, z);
    }
  }

  OctTree.Erase();

  int Planes = 1;
  if(indices >    2) Planes=2;
  if(indices >    4) Planes=4;
  if(indices >   16) Planes=8;
  if(indices >  256) Planes=16;
  if(indices >65536) Planes=32;
  if(Planes>=16) return NULL;

  APalette *Palette = BuildPalette(indices,8);
  if(Palette)
  {
    for(unsigned i=0; i<Palette->Size1D; i++)
    {
      Palette->setR(i, Arr[i].R * 0x11);
      Palette->setG(i, Arr[i].G * 0x11);
      Palette->setB(i, Arr[i].B * 0x11);
    }
  }
  return Palette;
}


///////////////////////////////////////////////////


Raster2DAbstract *Remap2Palette(const APalette *Palette, Raster2DAbstract *RasterRGB)
{
Raster2DAbstract *Ras;
  if(Palette==NULL || RasterRGB==NULL) return NULL;
  if(Palette->Size1D==0 || RasterRGB->Size1D==0 || RasterRGB->Size2D==0) return NULL;
  Raster1D_8BitRGB RasPalRGB;
  const uint8_t *PalRGB;
  if(Palette->GetPlanes()==24 && Palette->Channels()==3)
    PalRGB = (uint8_t*)Palette->Data1D;
  else
  {
    RasPalRGB.Set(*Palette);
    PalRGB = (uint8_t*)RasPalRGB.Data1D;
    if(PalRGB == NULL) return NULL;
  }

  int Planes = 1;
  if(Palette->Size1D >    2) Planes=2;
  if(Palette->Size1D >    4) Planes=4;
  if(Palette->Size1D >   16) Planes=8;
  if(Palette->Size1D >  256) Planes=16;
  if(Palette->Size1D >65536) Planes=32;
  if(Planes>=16) return NULL;

  Ras = CreateRaster2D(RasterRGB->Size1D,RasterRGB->Size2D, Planes);
  if(Ras==NULL) return NULL;

#pragma omp parallel
  {
    Raster1DAbstract * const RasRGBrow 
	  =  (RasterRGB->Channels()==3) ? CreateRaster1DRGB(0,RasterRGB->GetPlanes()/3) : CreateRaster1DRGBA(0,RasterRGB->GetPlanes()/4);
    RasRGBrow->Shadow = true;
    RasRGBrow->Size1D = RasterRGB->Size1D;
    Raster1DAbstract * const RasRow = CreateRaster1D(0, Planes);
    RasRow->Shadow = true;
    RasRow->Size1D = Ras->Size1D;

#pragma omp for ordered schedule(dynamic)
    for(UNS_OMP_IT y=0; y<RasterRGB->Size2D; y++)
    {
      RasRGBrow->Data1D = RasterRGB->GetRow(y);
      RasRow->Data1D = Ras->GetRow(y);
      for(unsigned x=0; x<RasterRGB->Size1D; x++)
      {
        unsigned BestIDX = 0;
        uint32_t BestDiff = 0xFFFFFFFF;
        RGBQuad RGB;
        RasRGBrow->Get(x,&RGB);

        for(unsigned c=0; c<Palette->Size1D; c++)
        {
          const uint32_t Diff = sqr((int32_t)RGB.R-(int32_t)PalRGB[3*c]) + sqr((int32_t)RGB.G-(int32_t)PalRGB[3*c+1]) + sqr((int32_t)RGB.B-(int32_t)PalRGB[3*c+2]);
          if(Diff < BestDiff)
            {BestIDX=c; BestDiff=Diff;}
        }
        RasRow->SetValue1D(x,BestIDX);
      }
    }

    delete RasRGBrow;
    delete RasRow;
  }

return Ras;
}


double GetCriteriaFit(const APalette *Palette, Raster2DAbstract *RasterRGB)
{
double SumCriteria = 0;
Raster1D_8BitRGB RasPalRGB;
const uint8_t *PalRGB;

  if(Palette==NULL || RasterRGB==NULL) return -1;
  if(Palette->Size1D==0 || RasterRGB->Size1D==0 || RasterRGB->Size2D==0) return -1;

  if(Palette->GetPlanes()==24 && Palette->Channels()==3)
    PalRGB = (uint8_t*)Palette->Data1D;
  else
  {
    RasPalRGB.Set(*Palette);
    PalRGB = (uint8_t*)RasPalRGB.Data1D;
    if(PalRGB == NULL) return -1;
  }

  #pragma omp parallel
  {
    Raster1DAbstract * const RasRGBrow 
	  = CreateRaster1DRGB(0,RasterRGB->GetPlanes());
    RasRGBrow->Shadow = true;
    RasRGBrow->Size1D = RasterRGB->Size1D;

  #pragma omp for ordered schedule(dynamic)
    for(UNS_OMP_IT y=0; y<RasterRGB->Size2D; y++)
    {
      RasRGBrow->Data1D = RasterRGB->GetRow(y);

    //
      for(unsigned x=0; x<RasterRGB->Size1D; x++)
      {      
        uint32_t BestDiff = 0xFFFFFFFF;
        RGBQuad RGB;

        RasRGBrow->Get(x,&RGB);

        for(unsigned c=0; c<Palette->Size1D; c++)
        {
          const uint32_t Diff = sqr((int32_t)RGB.R-(int32_t)PalRGB[3*c]) + sqr((int32_t)RGB.G-(int32_t)PalRGB[3*c+1]) + sqr((int32_t)RGB.B-(int32_t)PalRGB[3*c+2]);
          if(Diff < BestDiff)
            BestDiff=Diff;
        }
        SumCriteria += BestDiff;
      }    
    }

    delete RasRGBrow;
  }
  return SumCriteria;
}